package com.aloda.ai.impl;

import cn.hutool.core.util.ObjectUtil;
import com.aloda.ai.AiPlatform;
import com.aloda.ai.anno.AiPlatformConfig;
import com.aloda.constant.AiConstant;
import com.aloda.constant.ErrorCode;
import com.aloda.constant.ExceptionConstant;
import com.aloda.context.UserBaseContext;
import com.aloda.dto.AiGenerateResultDTO;
import com.aloda.dto.AiMessageDTO;
import com.aloda.dto.FlowableData;
import com.aloda.enums.AiPlatformEnum;
import com.aloda.exception.BaseException;
import com.volcengine.ark.runtime.model.completion.chat.*;
import com.volcengine.ark.runtime.service.ArkService;
import io.reactivex.Flowable;
import jakarta.annotation.Resource;
import lombok.extern.slf4j.Slf4j;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.stream.Collectors;

@Slf4j
@AiPlatformConfig(platformType = AiPlatformEnum.VOLCANO_AI)
public class VolcanoAiPlatform implements AiPlatform {

    @Resource
    private ArkService arkService;

    /**
     * 火山引擎引入的大模型名称枚举
     */
    public enum VolcanoModelEnum {
        DEEP_SEEK_V3("deepseek-v3-250324", 917),
        DOU_BAO_PRO_32K("doubao-1.5-pro-32k-250115", 420);

        private final String text;

        private final int value;

        VolcanoModelEnum(String text, int value) {
            this.text = text;
            this.value = value;
        }

        /**
         * 根据 value 获取枚举
         *
         * @param value
         * @return
         */
        public static VolcanoModelEnum getEnumByValue(Integer value) {
            if (ObjectUtil.isEmpty(value)) {
                return null;
            }
            for (VolcanoModelEnum anEnum : VolcanoModelEnum.values()) {
                if (anEnum.value == value) {
                    return anEnum;
                }
            }
            return null;
        }

        /**
         * 获取值列表
         *
         * @return
         */
        public static List<Integer> getValues() {
            return Arrays.stream(values()).map(item -> item.value).collect(Collectors.toList());
        }

        public int getValue() {
            return value;
        }

        public String getText() {
            return text;
        }
    }


    /**
     * 处理同步调用请求
     *
     * @param aiMessageDTO AI请求数据传输对象
     * @return AiGenerateResultDTO AI响应数据传输对象
     */
    @Override
    public AiGenerateResultDTO doSyncRequest(AiMessageDTO aiMessageDTO) {
        // 1.构建请求消息
        List<ChatMessage> messageList = getMessage(aiMessageDTO);

        // 2.选择模型
        VolcanoModelEnum modelEnum = VolcanoModelEnum.getEnumByValue(aiMessageDTO.getModel());
        if (modelEnum == null) {
            throw new BaseException(ExceptionConstant.MODEL_NOT_EXIST);
        }
        log.info("用户：{}开始调用AI服务:平台:{}, 模型:{}",
                UserBaseContext.getUser().getId(),
                aiMessageDTO.getAiPlatform().getText(),
                modelEnum.getText());
        // 3.构建请求
        ChatCompletionRequest chatCompletionRequest = ChatCompletionRequest.builder()
                .model(modelEnum.getText())
                .stream(Boolean.FALSE)
                .temperature((double)aiMessageDTO.getTemperature())
                .messages(messageList)
                .maxTokens(AiConstant.MAX_TOKENS)
                .build();

        // 4.发送请求并返回结果
        try {
            ChatCompletionResult chatCompletion = arkService.createChatCompletion(chatCompletionRequest);
            String choices = chatCompletion.getChoices().get(0).toString();
            long totalTokens = chatCompletion.getUsage().getTotalTokens();
            // 封装简单返回数据
            AiGenerateResultDTO aiGenerateResultDTO = new AiGenerateResultDTO();
            aiGenerateResultDTO.setResult(choices);
            aiGenerateResultDTO.setTokenNum(totalTokens);
            log.info("用户：{}完成调用AI服务:平台:{}, 模型:{},消耗tokens:{}",
                    UserBaseContext.getUser().getId(),
                    aiMessageDTO.getAiPlatform().getText(),
                    modelEnum.getText(),
                    totalTokens);
            return aiGenerateResultDTO;
        } catch (Exception e) {
            e.printStackTrace();
            throw new BaseException(ErrorCode.SYSTEM_ERROR);
        }
    }

    /**
     * 通用流式请求
     * @return
     */
    public Flowable<FlowableData> doStreamRequest(AiMessageDTO aiMessageDTO) {
        // 1.构建请求消息
        List<ChatMessage> messageList = getMessage(aiMessageDTO);
        // 2.选择模型
        VolcanoModelEnum modelEnum = VolcanoModelEnum.getEnumByValue(aiMessageDTO.getModel());
        if (modelEnum == null) {
            throw new BaseException(ExceptionConstant.MODEL_NOT_EXIST);
        }
        Long userId = UserBaseContext.getUser().getId();
        log.info("用户：{}开始调用AI服务:平台:{}, 模型:{}",
                userId,
                aiMessageDTO.getAiPlatform().getText(),
                modelEnum.getText());
        // 3.构建请求
        ChatCompletionRequest.ChatCompletionRequestStreamOptions chatCompletionRequestStreamOptions =
                new ChatCompletionRequest.ChatCompletionRequestStreamOptions(Boolean.TRUE);
        ChatCompletionRequest chatCompletionRequest = ChatCompletionRequest.builder()
                .model(modelEnum.getText())
                .stream(Boolean.TRUE)
                .streamOptions(chatCompletionRequestStreamOptions)
                .temperature((double)aiMessageDTO.getTemperature())
                .messages(messageList)
                .maxTokens(AiConstant.MAX_TOKENS)
                .build();
        // 4.发送请求并返回结果
        Flowable<ChatCompletionChunk> chatCompletionChunkFlowable = arkService.streamChatCompletion(chatCompletionRequest);
        return chatCompletionChunkFlowable.map(chunk -> {
            FlowableData flowableData = new FlowableData();
            flowableData.setMessage("");
            if (chunk.getUsage() != null) {
                flowableData.setTotalTokens(chunk.getUsage().getTotalTokens());
                log.info("用户：{}完成调用AI服务:平台:{}, 模型:{},消耗tokens:{}",
                        userId,
                        aiMessageDTO.getAiPlatform().getText(),
                        modelEnum.getText(),
                        chunk.getUsage().getTotalTokens());
            }
            if (chunk.getChoices().size() > 0) {
                flowableData.setMessage(chunk.getChoices().get(0).getMessage().getContent().toString());
            }
            return flowableData;
        });
    }

    private List<ChatMessage> getMessage(AiMessageDTO aiMessageDTO) {
        List<ChatMessage> messages = new ArrayList<>();
        ChatMessage systemMessage = ChatMessage
                .builder().role(ChatMessageRole.SYSTEM).content(aiMessageDTO.getSystemMessage()).build();
        ChatMessage userMessage = ChatMessage
                .builder().role(ChatMessageRole.USER).content(aiMessageDTO.getUserMessage()).build();
        messages.add(systemMessage);
        messages.add(userMessage);
        return messages;
    }
}
